
package com.spring.task.util;

import org.springframework.jdbc.core.BeanPropertyRowMapper;
import org.springframework.jdbc.core.namedparam.MapSqlParameterSource;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;


public class DBDataStream<T> {

	private final String sql;
	private final NamedParameterJdbcTemplate jdbcTemplate;
	private int pageSize;
	private int pageNo;
	private int productsCount;
	private int pageCount;
	private Class<T> clz;

	public DBDataStream(String sql, NamedParameterJdbcTemplate jdbcTemplate, int pageSize,
                        Class<T> clz) {
		this.sql = sql;
		this.jdbcTemplate = jdbcTemplate;
		this.pageSize = pageSize == 0 ? 50 : pageSize;
		this.clz = clz;
	}

	public Stream<T> getDataStream(MapSqlParameterSource param) {
		productsCount = getTotalSize(param);
		pageCount = getPageCount(productsCount);
		Iterator<T> iter = new Iterator<T>() {
			ConcurrentLinkedQueue<T> ts = new ConcurrentLinkedQueue<>();

			@Override
			public boolean hasNext() {
				if (!CollectionUtils.isEmpty(ts)) {
					return true;
				} else {
					ts.addAll(queryPageData(param));
					return (!CollectionUtils.isEmpty(ts));
				}
			}

			@Override
			public T next() {
				if (!CollectionUtils.isEmpty(ts) || hasNext()) {
					return ts.poll();
				} else {
					throw new NoSuchElementException();
				}
			}
		};
		return StreamSupport.stream(Spliterators.spliterator(iter, productsCount,
				Spliterator.ORDERED | Spliterator.NONNULL), true);
	}

	private int getTotalSize(MapSqlParameterSource param) {
		Integer count = jdbcTemplate.queryForObject(getCountSql(), param, Integer.class);
		return count == null ? 0 : count;
	}

	private int getPageCount(int total) {
		return (total + pageSize - 1) / pageSize;
	}

	/** 返回count的sql并自动移除order by */
	private String getCountSql() {
		String countSql = "";
		if (sql.indexOf(" group ") != -1 || sql.indexOf(" GROUP ") != -1) {
			countSql = "select count(1) from (" + sql + ")  r";
		} else {
			int posFrom = sql.indexOf(" from ");
			if (posFrom == -1) throw new IllegalArgumentException("SQL中不存在' from ', " + sql);
			countSql = "select count(*)" + sql.substring(posFrom);
		}

//		int posOrderBy = countSql.lastIndexOf(" order by ");
//		if (posOrderBy != -1 && countSql.indexOf(")", posOrderBy) == -1)
//			countSql = countSql.substring(0, posOrderBy);
		return countSql;
	}

	private List<T> queryPageData(MapSqlParameterSource param) {
		MapSqlParameterSource parameters = param;
		if(parameters == null){
			parameters = new MapSqlParameterSource();
		}
		if (productsCount > 0) {
			if (pageCount > pageNo) {
//				String limitSql = sql + " limit :offset, :pageSize";
				String limitSql = sql + " where rownum BETWEEN :offset and :pageSize";
				parameters.addValue("offset", pageNo * pageSize);
				parameters.addValue("pageSize", (pageNo +1 )*pageSize-1);
				pageNo++;
				return jdbcTemplate.query(limitSql, parameters, new BeanPropertyRowMapper<T>(clz));
			}
		}
		return Collections.emptyList();
	}
}
